#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <cublas.h>
#include <cuda.h>
#include <R.h>
//#include <cuseful.h>

// C wrapper for CUBLAS function
extern "C" {
void gpuMatMult(int tpA, int tpB, float *a, int rowsa, int colsa,
        float *b, int rowsb, int colsb, float *c)
{
    // init GPU data
    float *gpua, *gpub, *gpuc;

    // allocate memory
    cublasInit();
//    checkCublasError("gpuMatMult device initialization.");
    cublasAlloc(rowsa*colsa, sizeof(float), (void**)&gpua);
//    checkCublasError("gpuMatMult matrix A allocation.");
    cublasAlloc(rowsb*colsb, sizeof(float), (void**)&gpub);
//    checkCublasError("gpuMatMult matrix B allocation.");
    
    char opA = tpA ? 'T' : 'N';
    char opB = tpB ? 'T' : 'N';
    int rowsOpA = tpA ? colsa : rowsa;
    int colsOpA = tpA ? rowsa : colsa;
    int colsOpB = tpB ? rowsb : colsb;
    cublasAlloc(rowsOpA*colsOpB, sizeof(float), (void **)&gpuc);
//    checkCublasError("gpuMatMult matrix C allocation.");

    // set the data
    cublasSetVector(rowsa*colsa, sizeof(float), a, 1, gpua, 1);
    cublasSetVector(rowsb*colsb, sizeof(float), b, 1, gpub, 1);

    // make CUBLAS call
    cublasSgemm(opA, opB, rowsOpA, colsOpB, colsOpA, 1.0, gpua, rowsa, gpub, rowsb, 0.0, gpuc, rowsOpA);

    // get result
    cublasGetVector(rowsOpA*colsOpB, sizeof(float), gpuc, 1, c, 1);
//    checkCublasError("gpuMatMult read from device memory.");

    // free device memory
    cublasFree(gpua);
    cublasFree(gpub);
    cublasFree(gpuc);
//    cublasShutDown();
}
}


// R wrapper for C function
extern "C" {
void RgpuMatMult(int *tpA, int *tpB, float *a, int *rowsa, int *colsa,
    float *b, int *rowsb, int *colsb, float *result)
{
    gpuMatMult(*tpA, *tpB, a, *rowsa, *colsa, b, *rowsb, *colsb, result);
}
}
